package service

import (
	"context"
	"encoding/json"
	"fmt"
	chromago "github.com/amikos-tech/chroma-go"
	"github.com/gorilla/websocket"
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/ollama"
	"gorm.io/gorm"
	"lib/structs/lms"
	"lib/structs/tables"
	"lms/global"
	"sync"
)

var KnowledgeService = new(knowledgeService)

type knowledgeService struct {
}

var promMap = map[int64]string{
	1: "你是一个信息助手，请根据提示信息回答用户的问题。提示信息：",
}

var chatStory = sync.Map{}

// Chat 基于ws方式和页面交互
func (s knowledgeService) Chat(uid int64, resp *websocket.Conn, in lms.StartChat) (err error) {
	// 1.创建llm模型对象
	llm, err := ollama.New(
		ollama.WithModel(in.Model),
		ollama.WithServerURL(global.LmsConf.RAGConfig.OllamaUrl),
	)
	if err != nil {
		return err
	}
	// 2.查询引用文档
	var docs []lms.DocRefer
	if in.KnowledgeId != 0 {
		docs, err = CoursewareService.FindDoc(in.KnowledgeId, in.Query, in.ReferDocNum)
		if err != nil {
			return
		}
	}
	// 3.构建提示词
	msgs, err := s.getPromMsg(uid, docs, in)
	if err != nil {
		return
	}
	ctx := context.Background()
	// 4.开始回答
	completion, err := llm.GenerateContent(ctx,
		msgs,
		llms.WithTemperature(in.Temperature),
		llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
			err = resp.WriteMessage(websocket.TextMessage, chunk)
			if err != nil {
				return err
			}
			return nil
		}),
	)
	// 5.回答完毕后将文档引用信息发送给页面
	if len(docs) > 0 {
		bs, err := json.Marshal(docs)
		if err != nil {
			return err
		}
		// 文档json分隔符
		dbs := []byte("{doc=>}")
		dbs = append(dbs, bs...)
		err = resp.WriteMessage(websocket.TextMessage, dbs)
		if err != nil {
			return err
		}
	}
	// 6.对聊天记录进行保存
	msgs = append(msgs, llms.MessageContent{
		Role:  llms.ChatMessageTypeAI,
		Parts: []llms.ContentPart{llms.TextContent{Text: completion.Choices[0].Content}},
	})
	s.storeChat(uid, msgs)
	return nil
}

// getPromMsg 获取提示词
func (s knowledgeService) getPromMsg(uid int64, docs []lms.DocRefer, in lms.StartChat) (li []llms.MessageContent, err error) {
	li = s.getChat(uid)
	// 不使用知识库的情况下不用加载提示词和文档
	if in.KnowledgeId != 0 {
		// 装载提示词
		if len(li) == 0 {
			li = append(li, llms.MessageContent{
				Role:  llms.ChatMessageTypeSystem,
				Parts: []llms.ContentPart{llms.TextContent{Text: promMap[in.Prom]}},
			})
		}
		// 装载提示文档
		for _, v := range docs {
			li = append(li, llms.MessageContent{
				Role:  llms.ChatMessageTypeSystem,
				Parts: []llms.ContentPart{llms.TextContent{Text: v.Content}},
			})
		}
	}
	// 装载用户的问题
	li = append(li, llms.MessageContent{
		Role:  llms.ChatMessageTypeHuman,
		Parts: []llms.ContentPart{llms.TextContent{Text: in.Query}},
	})
	return
}

// storeChat 保存聊天记录，注意只会保存最近的20条聊天记录记录,固定为第一条提示词+后面的19条数据
func (s knowledgeService) storeChat(uid int64, li []llms.MessageContent) {
	var rli []llms.MessageContent
	if len(li) <= 20 {
		rli = li
	} else {
		rli = make([]llms.MessageContent, 0)
		// 第一条是提示词，在没有添加自定义提示词切换功能前这条信息不能丢。
		one := li[0]
		last := li[len(li)-19:]
		rli = append(rli, one)
		rli = append(rli, last...)
	}
	chatStory.Store(uid, rli)
}

// getChat 获取历史聊天记录,包含第一条提示词和最后的8条信息，合计9条
func (s knowledgeService) getChat(uid int64) (li []llms.MessageContent) {
	v, ok := chatStory.Load(uid)
	if !ok {
		li = make([]llms.MessageContent, 0)
		return
	}
	li = v.([]llms.MessageContent)
	if len(li) < 10 {
		return li
	}
	// 第一条是提示词，在没有添加自定义提示词切换功能前这条信息不能丢。
	one := li[0]
	// 额外在取后面的8条信息，+提示词1，+用户刚提出的问题，一共10条
	last := li[len(li)-8:]
	rli := make([]llms.MessageContent, 0)
	rli = append(rli, one)
	rli = append(rli, last...)
	return rli
}

// DeleteById 删除知识库，同步删除向量库
func (s knowledgeService) DeleteById(id int64) (err error) {
	err = global.DB.Transaction(func(tx *gorm.DB) error {
		// 1.删除知识库
		tx.Delete(&tables.Knowledge{Id: id})
		// 2.删除知识库和课件映射
		tx.Exec("delete from lms_map_knowledge_doc where knowledge_id = ?", id)
		// 3.删除向量库
		err = s.deleteChromeDBCollection(fmt.Sprintf("%d", id))
		return err

	})
	return
}

// deleteChromeDBCollection 删除chroma数据集
func (s knowledgeService) deleteChromeDBCollection(collName string) (err error) {
	client, err := chromago.NewClient(global.LmsConf.RAGConfig.ChromaDBUrl)
	if err != nil {
		return err
	}
	_, err = client.DeleteCollection(context.Background(), collName)
	return
}
